package io.gatling.extensions.mongodb.protocol

import java.util.concurrent.TimeUnit

import com.mongodb.MongoClientSettings
import com.mongodb.connection.ConnectionPoolSettings
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.gatling.core.{CoreComponents, protocol}
import org.mongodb.scala.connection.ConnectionPoolSettings.Builder
import org.mongodb.scala.{ConnectionString, MongoClient}

case class MongoDBProtocol(url: String, pollSize: Int = 100, waitTime: Int = 10) extends Protocol {
  type Components = MongoDBComponents
}

object MongoDBProtocol {

  def apply(url: String, pollSize: Int = 100, waitTime: Int = 10): MongoDBProtocol = new MongoDBProtocol(url, pollSize, waitTime)

  val mongoDBProtocolKey: ProtocolKey[MongoDBProtocol, MongoDBComponents] = new ProtocolKey[MongoDBProtocol, MongoDBComponents] {

    override def protocolClass: Class[protocol.Protocol] = classOf[MongoDBProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): MongoDBProtocol = MongoDBProtocol("mongodb://root:auto_test%40123@192.168.106.147:27017/admin")

    override def newComponents(coreComponents: CoreComponents): MongoDBProtocol => MongoDBComponents = {
      mongoDBProtocol => MongoDBComponents(mongoDBProtocol)
    }
  }
}

case class MongoDBComponents(mongoDBProtocol: MongoDBProtocol) extends ProtocolComponents {

  System.setProperty("org.mongodb.async.type", "netty")

  // 连接池配置
  val pollSettingBuilder: ConnectionPoolSettings.Builder = ConnectionPoolSettings.builder()
    .minSize(mongoDBProtocol.pollSize)
    .maxSize(mongoDBProtocol.pollSize)
    .maxWaitTime(mongoDBProtocol.waitTime, TimeUnit.SECONDS)

  // 客户端配置
  val clientSettingbuilder: MongoClientSettings.Builder = MongoClientSettings.builder()
    .applyConnectionString(ConnectionString(mongoDBProtocol.url))
    .applyToConnectionPoolSettings((builder: Builder) => builder.applySettings(pollSettingBuilder.build()).build())

  // 创建客户端
  val client: MongoClient = MongoClient(clientSettingbuilder.build())

  override def onStart: Session => Session = session => session

  override def onExit: Session => Unit = _ => ???
}

case class MongoDBProtocolBuilder(url: String) {

  private var pollSize: Int = 100
  private var waitTime: Int = 10

  def pollSize(size: Int): MongoDBProtocolBuilder = {
    this.pollSize = size
    this
  }

  def waitTime(time: Int): MongoDBProtocolBuilder = {
    this.waitTime = time
    this
  }

  def build: MongoDBProtocol = MongoDBProtocol(url, pollSize, waitTime)
}

object MongoDBProtocolBuilder {
  implicit def toMongoProtocol(builder: MongoDBProtocolBuilder): MongoDBProtocol = builder.build
}
